# 初始库
import sys
sys.path.append('..')
import numpy as np
import pandas as pd
import empyrical as ep
import scipy.stats as st
import datetime
from DE_algorithm import *
from tqdm import tqdm_notebook
from typing import (Tuple,List)
from dateutil.parser import parse
import matplotlib as mpl
import matplotlib.pyplot as plt
plt.style.use('ggplot')

# 设置字体 用来正常显示中文标签
#plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['font.family']='serif'
# 用来正常显示负号
plt.rcParams['axes.unicode_minus'] = False
# 图表主题
plt.style.use('seaborn')


# 构造时间区间
def build_periods(start: str, end: str, freq: str, offsetN: int = None) -> Tuple[list, list]:
    '''构造区间'''

    # 半年度要特殊一点
    if freq == 'half':

        periods = [i.strftime(
            '%Y-%m-%d') for i in pd.date_range(start, end, freq='Q') if i.month in [6, 12]]

        end_range = periods[1:] + [end]
        start_range = periods

    else:

        periods = [i.strftime('%Y-%m-%d')
                   for i in pd.date_range(start, end, freq=freq)]

        start_range = periods[:-1]
        end_range = periods[1:]

    if offsetN:

        start_range = [tdaysoffset(i, offsetN) for i in start_range]
        end_range = [tdaysoffset(i, offsetN) for i in end_range]

        return (start_range, end_range)

    else:

        return (start_range, end_range)


# # offset day
# def tdaysoffset(end_date: str, count: int) -> datetime.date:
#     '''
#     end_date:为基准日期
#     count:为正则后推，负为前推
#     -----------
#     return datetime.date
#     '''
#
#     trade_date = get_trade_days(end_date=end_date, count=1)[0]
#
#     if count > 0:
#         # 将end_date转为交易日
#
#         trade_cal = get_all_trade_days().tolist()
#
#         trade_idx = trade_cal.index(trade_date)
#
#         return trade_cal[trade_idx + count]
#
#     elif count < 0:
#
#         return get_trade_days(end_date=trade_date, count=abs(count))[0]
#
#     else:
#
#         raise ValueError('别闹！')


# 风险指标
def Strategy_performance(return_df: pd.DataFrame, periods='monthly') -> pd.DataFrame:
    '''计算风险指标 默认为月度:月度调仓'''

    ser: pd.DataFrame = pd.DataFrame()
    ser['年化收益率'] = ep.annual_return(return_df, period=periods)
    ser['波动率'] = return_df.apply(lambda x: ep.annual_volatility(x, period=periods))
    ser['夏普'] = return_df.apply(ep.sharpe_ratio, period=periods)
    ser['最大回撤'] = return_df.apply(lambda x: ep.max_drawdown(x))

    if 'benchmark' in return_df.columns:
        select_col = [col for col in return_df.columns if col != 'benchmark']

        ser['IR'] = return_df[select_col].apply(
            lambda x: information_ratio(x, return_df['benchmark']))
        ser['Alpha'] = return_df[select_col].apply(
            lambda x: ep.alpha(x, return_df['benchmark'], period=periods))

    return ser.T


def information_ratio(returns, factor_returns):
    """
    Determines the Information ratio of a strategy.

    Parameters
    ----------
    returns : :py:class:`pandas.Series` or pd.DataFrame
        Daily returns of the strategy, noncumulative.
        See full explanation in :func:`~empyrical.stats.cum_returns`.
    factor_returns: :class:`float` / :py:class:`pandas.Series`
        Benchmark return to compare returns against.

    Returns
    -------
    :class:`float`
        The information ratio.

    Note
    -----
    See https://en.wikipedia.org/wiki/information_ratio for more details.

    """
    if len(returns) < 2:
        return np.nan

    active_return = _adjust_returns(returns, factor_returns)
    tracking_error = np.std(active_return, ddof=1)
    if np.isnan(tracking_error):
        return 0.0
    if tracking_error == 0:
        return np.nan
    return np.mean(active_return) / tracking_error


def _adjust_returns(returns, adjustment_factor):
    """
    Returns a new :py:class:`pandas.Series` adjusted by adjustment_factor.
    Optimizes for the case of adjustment_factor being 0.

    Parameters
    ----------
    returns : :py:class:`pandas.Series`
    adjustment_factor : :py:class:`pandas.Series` / :class:`float`

    Returns
    -------
    :py:class:`pandas.Series`
    """
    if isinstance(adjustment_factor, (float, int)) and adjustment_factor == 0:
        return returns.copy()
    return returns - adjustment_factor


# 用于回测
def BackTesting(timeRange: List, weight: dict) -> pd.DataFrame:
    '''
    timeRange:时间区间
    wieght:权重key-trade,values:pd.Series
    ----------------
        return 等权、优化权重、基准的收益率序列
    '''

    begin = min(timeRange[0])
    end = max(timeRange[1])
    idx = pd.to_datetime(get_trade_days(begin, end))
    # 缓存容器
    ret_df = pd.DataFrame(data=np.zeros((len(idx), 2)),
                          index=idx,
                          columns=['opt_ret', 'equal_ret'])

    # 模拟回测
    ## 季度:beginDt日买入,至nextDt轮动
    for beginDt, nextDt in tqdm_notebook(zip(*timeRange), total=len(timeRange[0])):
        # 权重优化基于beginDt
        weight_ser = weight[beginDt]  # 读取权重
        code = weight_ser.index.tolist()

        # 持仓从beginDt日至nextDt日
        price = get_price(code, beginDt, tdaysoffset(nextDt, 1), fields='close', panel=False)
        close = pd.pivot_table(price, index='time', columns='code', values='close')
        pct_df = close.pct_change().shift(-1).iloc[:-1]  # 收益滞后一期 防止未来数据

        ret_df.loc[pct_df.index, 'opt_ret'] = pct_df @ weight_ser  # 计算组合权重
        ret_df.loc[pct_df.index, 'equal_ret'] = pct_df.mean(axis=1)  # 等权

    benchmark = get_price('000300.XSHG', startDt, endDt, fields='close')
    benchmark = benchmark.pct_change().reindex(ret_df.index)
    ret_df['benchmark'] = benchmark
    return ret_df


# #####  超参数 寻找在de进化算法下的最优初始参数
from sklearn.base import TransformerMixin, BaseEstimator
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV
from typing import (Callable)
'''使用sklearn接口用于超参'''


class DEPortfolioOpt(BaseEstimator):
    def __init__(self, func: Callable, size_pop: int, max_iter: int, F: float, proub_mut: float) -> None:
        self.func = func  # 目标函数
        self.size_pop = size_pop  # 种群大小
        self.max_iter = max_iter  # 迭代次数
        self.F = F  # 变异系数
        self.proub_mut = proub_mut  # 变异概率

    def fit(self, returns: pd.DataFrame) -> np.array:
        '''获取优化后的权重'''

        self.de = DE(func=lambda w: self.func(returns, w),
                     n_dim=returns.shape[1],
                     size_pop=self.size_pop,
                     max_iter=self.max_iter,
                     F=self.F,
                     prob_mut=self.proub_mut,
                     lb=[0],
                     ub=[1],
                     constraint_eq=[lambda x: 1 - np.sum(x)])

        w, self.maxDrawdown = self.de.run()

        return pd.Series(w, index=returns.columns)

    def predict(self, returns) -> float:
        '''获取权重后的收益率'''
        w = self.fit(returns)
        return returns @ w

    def score(self, returns) -> float:
        '''评分根据优化后的回撤决定'''
        return -self.maxDrawdown

